// Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashMap;

use fixedbitset::FixedBitSet;
use petgraph::{graph::{EdgeIndex, NodeIndex}, visit::EdgeRef, Graph};
use strum::EnumCount;

use crate::{config::{pkg_config::PkgConfigLoadError, system::{toolchain::{SubToolchainKind, ToolchainConfig}, OverlayConfig}, PkgDepConfig, SpecsConfig}, graph::build_dep_graph::ToolchainRanges, Arch, Pkg, PkgSpec};

use super::{build_dep_graph::{DepEdge, DepGraph, DepNode}, DepKind};

pub struct PkgConfigResolver<'a> {
    pub dep_graph: DepGraph,
    unresolved: Vec<(DepInfo, PkgSpec)>,
    toolchain: &'a ToolchainConfig,
}

struct DepInfo {
    dependent: NodeIndex,
    edge: DepEdge,
}

#[derive(Debug, thiserror::Error)]
pub enum ResolveError {
    #[error(transparent)]
    PkgConfigLoadFailed(#[from] PkgConfigLoadError),
    #[error("loop dependency of pkg '{0}'")]
    LoopDependency(String),
    #[error("pkg '{0}' has no feature '{1}'")]
    FeatureNotFound(String, String),
    #[error("toolchain '{1}' required by pkg '{0}' not found")]
    ToolchainNotFound(String, SubToolchainKind),
}

impl<'a> PkgConfigResolver<'a> {
    pub fn new(toolchain: &'a ToolchainConfig) -> Self {
        Self {
            dep_graph: DepGraph::default(),
            unresolved: Vec::new(),
            toolchain,
        }
    }

    pub async fn build_graph<I>(
        &mut self,
        specs: I,
    ) -> Result<(), ResolveError>
    where
        I: Iterator<Item = PkgSpec>,
    {
        self.preload_toolchain().await?;
        self.preload_specs(specs).await?;

        while let Some((dep_info, spec)) = self.unresolved.pop() {
            let id = self.resolve_spec(dep_info, spec).await?;
            if id.index() + 1 == self.dep_graph.node_count() {
                self.dep_graph.add_toolchain_deps(id)?;
            }
        }

        self.dep_graph.add_deploy_runtime_deps();

        Ok(())
    }

    async fn preload_toolchain(&mut self) -> Result<(), ResolveError> {
        async fn add_toolchain_nodes(
            dep_graph: &mut DepGraph,
            toolchain_node_map: &mut HashMap<String, NodeIndex>,
            spec_config: &SpecsConfig,
            overlay: Option<&OverlayConfig>,
        ) -> Result<(), ResolveError> {
            for (k, v) in spec_config.0.iter() {
                let spec = v.to_pkg_spec(k);
                debug_assert!(spec.features.is_sorted());
                let id = dep_graph.add_node(spec, overlay).await?;
                toolchain_node_map.insert(k.clone(), id);
            }
            Ok(())
        }
        
        let mut toolchain_node_map = HashMap::new();

        for (kind, subtoolchain) in self.toolchain.sub_toolchains.iter() {
            let mut ranges = ToolchainRanges::default();

            let start = self.dep_graph.node_count();
            add_toolchain_nodes(
                &mut self.dep_graph,
                &mut toolchain_node_map,
                &subtoolchain.packages,
                self.toolchain.overlay.as_ref(),
            ).await?;
            ranges.host_toolchain_range = start..self.dep_graph.node_count();

            for (&arch, spec_config) in subtoolchain.cross_packages.iter() {
                let start = self.dep_graph.node_count();
                add_toolchain_nodes(
                    &mut self.dep_graph,
                    &mut toolchain_node_map,
                    spec_config,
                    self.toolchain.overlay.as_ref(),
                ).await?;
                ranges.cross_toolchain_ranges.insert(arch, start..self.dep_graph.node_count());
            }
            
            self.dep_graph.id_map.subtoolchain_ranges.insert(*kind, ranges);
        }

        let start = self.dep_graph.node_count();
        add_toolchain_nodes(
            &mut self.dep_graph,
            &mut toolchain_node_map,
            &self.toolchain.temp_packages,
            self.toolchain.overlay.as_ref(),
        ).await?;
        self.dep_graph.id_map.temp_toolchain_range = start..self.dep_graph.node_count();

        for extra_dep_edge in self.toolchain.extra_dep_edges.iter() {
            let source = toolchain_node_map.get(&extra_dep_edge.source).unwrap();
            let target = toolchain_node_map.get(&extra_dep_edge.target).unwrap();
            DepGraph::add_toolchain_dep_edge(
                &mut self.dep_graph.edge_map,
                &mut self.dep_graph.graph,
                *source,
                *target,
                extra_dep_edge.kind,
            );
        }

        self.dep_graph.id_map.preload_count = self.dep_graph.node_count();
        for id in 0..self.dep_graph.id_map.preload_count {
            self.add_depends(NodeIndex::new(id))?;
        }
        
        while let Some((dep_info, spec)) = self.unresolved.pop() {
            self.resolve_spec(dep_info, spec).await?;
        }

        Ok(())
    }

    async fn preload_specs<I>(&mut self, specs: I) -> Result<(), ResolveError>
    where
        I: Iterator<Item = PkgSpec>,
    {
        let start = self.dep_graph.node_count();
        self.dep_graph.id_map.toolchain_specs.grow(start);
        for spec in specs {
            debug_assert!(spec.features.is_sorted());
            let id = self.dep_graph.add_node(spec, self.toolchain.overlay.as_ref()).await?;
            if id.index() + 1 == self.dep_graph.node_count() {
                self.dep_graph.add_toolchain_deps(id)?;
            } else {
                self.dep_graph.id_map.toolchain_specs.put(id.index());
            }
        }
        self.dep_graph.id_map.preload_count = self.dep_graph.node_count();
        self.dep_graph.id_map.non_toolchain_specs_range = start..self.dep_graph.id_map.preload_count;

        for id in start..self.dep_graph.id_map.preload_count {
            self.add_depends(NodeIndex::new(id))?;
        }
        
        Ok(())
    }

    async fn resolve_spec(&mut self, dep_info: DepInfo, spec: PkgSpec) -> Result<NodeIndex, ResolveError> {
        let id = self.dep_graph.add_node(spec, self.toolchain.overlay.as_ref()).await?;
        if id.index() + 1 == self.dep_graph.node_count() {
            self.add_depends(id)?;
        }
        
        self.dep_graph.add_edge(id, dep_info.dependent, dep_info.edge);
        Ok(id)
    }

    fn add_depends(&mut self, dependent: NodeIndex) -> Result<(), ResolveError> {
        let node = self.dep_graph.node(dependent).unwrap();
        for (name, dep_config) in node.config.depends.non_optional_deps() {
            let spec = self.resolve_dep(name, dep_config, node.spec.arch);
            self.unresolved.push((DepInfo::new(dependent, DepKind::Build, dep_config), spec));
        }
        for (name, dep_config) in node.config.host_depends.non_optional_deps() {
            let spec = self.resolve_dep(name, dep_config, Arch::host());
            self.unresolved.push((DepInfo::new(dependent, DepKind::HostBuild, dep_config), spec));
        }
        for (name, dep_config) in node.config.runtime_depends.non_optional_deps() {
            let spec = self.resolve_dep(name, dep_config, node.spec.arch);
            self.unresolved.push((DepInfo::new(dependent, DepKind::Runtime, dep_config), spec));
        }

        for feature in node.spec_features() {
            let Some(feature_config) = node.config.features.configs.get(feature) else {
                return Err(ResolveError::FeatureNotFound(node.spec.pkg.to_string(), feature.clone()));
            };

            for (dep_config, name) in feature_config.optional_depends(&node.config.depends) {
                let spec = self.resolve_dep(name, dep_config, node.spec.arch);
                self.unresolved.push((DepInfo::new(dependent, DepKind::Build, dep_config), spec));
            }
            for (dep_config, name) in feature_config.optional_host_depends(&node.config.host_depends) {
                let spec = self.resolve_dep(name, dep_config, Arch::host());
                self.unresolved.push((DepInfo::new(dependent, DepKind::HostBuild, dep_config), spec));
            }
            for (dep_config, name) in feature_config.optional_runtime_depends(&node.config.runtime_depends) {
                let spec = self.resolve_dep(name, dep_config, node.spec.arch);
                self.unresolved.push((DepInfo::new(dependent, DepKind::Runtime, dep_config), spec));
            }
        }
        Ok(())
    }

    fn resolve_dep(
        &self,
        name: &str,
        dep_config: &PkgDepConfig,
        arch: Arch
    ) -> PkgSpec {
        let pkgname = if let Some(name) = dep_config.name.as_ref() {
            name.as_str()
        } else {
            name
        };
        
        let mut features = dep_config.features.clone();
        features.sort();

        for id in 0..self.dep_graph.id_map.preload_count {
            let node = self.dep_graph.node(NodeIndex::new(id)).unwrap();
            if node.spec.arch == arch
                && node.spec.pkg.name == pkgname
                && dep_config.req.matches(&node.spec.pkg.version)
                && Self::match_features(&features, node.spec_features()) {
                return node.spec.clone();
            }
        }

        PkgSpec {
            pkg: Pkg {
                name: pkgname.to_string(),
                version: dep_config.default.clone(),
            },
            features,
            arch,
        }
    }

    fn match_features(features: &[String], spec_features: &[String]) -> bool {
        for feature in features {
            if !spec_features.contains(feature) {
                return false;
            }
        }
        true
    }
}

impl DepInfo {
    fn new(dependent: NodeIndex, kind: DepKind, config: &PkgDepConfig) -> Self {
        Self {
            dependent,
            edge: DepEdge {
                kind,
                config: Some(config.into()),
            },
        }
    }
}

impl DepGraph {
    pub fn verify(&mut self, deploy_only: bool) -> Result<Vec<NodeIndex>, ResolveError> {
        let sorted_ids = self.toposort()
            .map_err(|err| {
                ResolveError::LoopDependency(
                    self.node(err.node_id()).unwrap().spec.pkg.to_string()
                )
            })?;

        if deploy_only {
            // for (name, version_req) in config.conflicts.iter() {
                // TODO: check conflicts in deployed
            // }
        }

        Ok(sorted_ids)
    }

    fn add_toolchain_deps(&mut self, id: NodeIndex) -> Result<(), ResolveError> {
        let node = self.node(id).unwrap();
        let arch = node.spec.arch;
        let need_cross = !arch.is_host()
            && !node.config.build_options().is_some_and(|o| o.dont_need_cross_compile);
        
        let mut active_subtoolchains = FixedBitSet::with_capacity(SubToolchainKind::COUNT);
        for kind in node.config.toolchains.iter() {
            if !self.id_map.subtoolchain_ranges.contains_key(kind) {
                return Err(ResolveError::ToolchainNotFound(node.spec.pkg.to_string(), *kind));
            }
            active_subtoolchains.put(*kind as usize);
        }

        for kind in active_subtoolchains.ones() {
            let kind = SubToolchainKind::try_from(kind as u8).unwrap();
            let ranges = self.id_map.subtoolchain_ranges.get(&kind).unwrap();
            for toolchain_dep_id in ranges.host_toolchain_range.clone().map(NodeIndex::new) {
                let toolchain_dep_arch = self.node(toolchain_dep_id).unwrap().spec.arch;
                Self::add_toolchain_dep(
                    &mut self.edge_map,
                    &mut self.graph,
                    toolchain_dep_id,
                    id,
                    toolchain_dep_arch,
                    arch,
                );
            }
            if need_cross {
                if let Some(range) = ranges.cross_toolchain_ranges.get(&arch) {
                    for cross_toolchain_dep_id in range.clone().map(NodeIndex::new) {
                        let cross_toolchain_dep_arch = self.node(cross_toolchain_dep_id).unwrap().spec.arch;
                        Self::add_toolchain_dep(
                            &mut self.edge_map,
                            &mut self.graph,
                            cross_toolchain_dep_id,
                            id,
                            cross_toolchain_dep_arch,
                            arch,
                        );
                    }
                }
            }
        }
        Ok(())
    }

    fn add_toolchain_dep(
        edge_map: &mut HashMap<(NodeIndex, NodeIndex, DepKind), EdgeIndex>,
        graph: &mut Graph<DepNode, DepEdge>,
        toolchain_dep_id: NodeIndex,
        id: NodeIndex,
        toolchain_dep_arch: Arch,
        arch: Arch,
    ) {
        if arch != toolchain_dep_arch {
            if toolchain_dep_arch == Arch::host() {
                Self::add_toolchain_dep_edge(edge_map, graph, toolchain_dep_id, id, DepKind::HostBuild);
            } else {
                Self::add_toolchain_dep_edge(edge_map, graph, toolchain_dep_id, id, DepKind::Build);
            }
        } else {
            Self::add_toolchain_dep_edge(edge_map, graph, toolchain_dep_id, id, DepKind::Build);
            Self::add_toolchain_dep_edge(edge_map, graph, toolchain_dep_id, id, DepKind::HostBuild);
            Self::add_toolchain_dep_edge(edge_map, graph, toolchain_dep_id, id, DepKind::Runtime);
        }
    }

    fn add_toolchain_dep_edge(
        edge_map: &mut HashMap<(NodeIndex, NodeIndex, DepKind), EdgeIndex>,
        graph: &mut Graph<DepNode, DepEdge>,
        toolchain_dep_id: NodeIndex,
        id: NodeIndex,
        kind: DepKind,
    ) -> EdgeIndex {
        Self::add_edge_inner(
            edge_map,
            graph,
            toolchain_dep_id, id, DepEdge {
                kind,
                config: None,
            }
        )
    }

    fn add_deploy_runtime_deps(&mut self) {
        let mut set = self.visit_map();
        let mut stack = Vec::new();
        for id in self.id_map.specs() {
            for edge in self.runtime_dep_edges(id) {
                let dep_id = edge.source().index();
                if !set.contains(dep_id) {
                    set.set(dep_id, true);
                    stack.push(dep_id);
                }
            }
        }
        while let Some(current_id) = stack.pop() {
            for edge in self.runtime_dep_edges(NodeIndex::new(current_id)) {
                let dep_id = edge.source().index();
                if !set.contains(dep_id) {
                    set.set(dep_id, true);
                    stack.push(dep_id);
                }
            }
        }
        self.id_map.deploy_runtime_depend_specs = set;
    }
}
